#! /usr/bin/perl
#
#    $Id$
#
# Copyright (C) 2008 by Peder Stray <peder.stray@gmail.com>
#

use strict;
use Pod::Usage;
use Getopt::Long qw(:config bundling auto_version auto_help);
use Text::Balanced qw(:ALL);
use Data::Dumper;

our($VERSION) = '$Revision$' =~ / (\d+) /;

=head1 NAME

mkfktriggers - Generate triggers to enforce foreign keys for sqlite

=head1 SYNOPSIS

mkfktriggers [B<--strict>] [B<--prefix>=I<PREFIX>]
[B<-o>E<nbsp>F<output>|B<--output>=F<output>] [F<input>E<nbsp>...]

=head1 DESCRIPTION

This script generates triggers for SQLite to enforce foreign key
constraints defined in table definitions.  It makes triggers to check
integrity, and it also handles cascading if specified in the foreign key
definition.

=head1 OPTIONS

=over

=cut

my @options;

# Strict checking of trigger definitions.
my $opt_strict;
push @options, 'strict!' => \$opt_strict;

=item B<--strict>

Do strick checking of foreign key definitions.  Raise an error if the
definition of the target table and columns of a foreign key hasn't been
seen.  This option can also be specified in the source SQL-file with the
following comment:

    --#strict

=cut

# Value to prefix all trigger names
my $opt_prefix;
push @options, 'prefix=s' => \$opt_prefix;

=pod 

=item B<--prefix> I<prefix>

Specify a prefix for the name of the generated triggers.  This is
usefull when you want to specify your own triggers too, and want control
over when each trigger for the same operation is fired.  Typically
triggers for the same operation are fired in alfanumeric ordering by the
name of the trigger.  This option can also be specified in the source
SQL-file in the following way:

    --#prefix t50

This will prefix all triggers with "t50_" so you can name your other
triggers with either lower ot higher numbers in the prefix to get them
to be run either before or after the generated triggers.  A trigger you
want to run before an insert would probably need to be run before the
generated trigger, so you should then name it "t40_foo".

=cut

# outpuf if not STDOUT
my $opt_output;
push @options, 'output|o=s' => \$opt_output;

=item B<-o>, B<--output> F<file>

Specify the output file of the generated SQL.  If not specified,
standard output is used.

=cut

# parse commandline
GetOptions(@options) or pod2usage();

unless (@ARGV) {
    pod2usage();
}

=back

=cut

# read in the whole SQL data
my $data = do { local $/ = undef; <>; };

# find any defaults in special comments
while ($data =~ /^--#(\w+)(?:[ \t]+(.+))?[ \t]*$/mg) {
    my($key,$val) = ($1,$2);
    for ($val) {
	s/\s+$//;
	s/^\s+//;
    }
    if ($key eq 'prefix') {
	$opt_prefix = $val unless defined $opt_prefix;
    }
    elsif ($key eq 'strict') {
	$opt_strict = 1 unless defined $opt_strict;
    }
}

if ($opt_prefix) {
    $opt_prefix =~ s/_*$/_/;
}

if ($opt_prefix =~ /\s/) {
    die "prefix ($opt_prefix) is not a single word, it should be a valid SQL identifier\n";
}

# kill all comments
$data =~ s,/\*[^*]*\*+([^/*][^*]*\*+)*/
           |--[^\n]*
	   |("(\\.|[^"\\])*"
	       |'(\\.|[^'\\])*'
	       |.[^/"'\\-]*),defined $2 ? $2 : "",gsex;

my @tokens;
my %tables;

while ($data =~
       /create\s+table(?:\s+if\s+not\s+exists)?\s+(\S+)\s*\((.*?)\)\s*;/imsg) {
    my($table,$coldefs) = ($1, $2);
    
    my %tdata = (
		 name => $table,
		);

    @tokens = extract_multiple($coldefs,
			       [ sub { extract_bracketed($_[0], q{()}) },
				 sub { extract_delimited($_[0], q{"'`}) },
				 qr/[+-]?\w+/,
				 qr/,/,
			       ],
			       undef, # split as much as we can
			       1, # ignore anything else
			      );
    
    while (@tokens) {
	my %data;
	my @utok;

	if (match('CONSTRAINT')) {
	    $data{constraint} = token(",") or next;
	}

	if (match('UNIQUE')) {
	    # UNIQUE ( column_name [, ... ] ) 
	    next;
	}
	elsif (match('PRIMARY', 'KEY')) {
	    # PRIMARY KEY ( column_name [, ... ] )
	    if ($tokens[0] = /^\((.*)\)$/) {
		my @cols = split ",", $1;
		token();
		for (@cols) {
		    s/^\s+//;
		    s/\s+$//;
		}
		if (@cols) {
		    $tdata{pk}{name} = $data{constraint};
		    $tdata{pk}{cols} = \@cols;
		}
	    }
	    next;
	}
	elsif (match('CHECK')) {
	    # CHECK ( expression )
	    next;
	}
	elsif (match('FOREIGN', 'KEY')) {
	    # FOREIGN KEY ( column_name [, ... ] )

	    my @srccols;
	    if ($tokens[0] =~ /^\((.*)\)$/) {
		@srccols = split ",", $1;
		token();
		for (@srccols) {
		    s/^\s+//;
		    s/\s+$//;
		}
	    }
	    else {
		next;
	    }

	    # REFERENCES reftable [ ( refcolumn [, ... ] ) ]
	    match('REFERENCES') or next;

	    my $fk = parse_refdef($table, \@srccols, $data{constraint})
	      or next;

	    push @{$tdata{fk}}, $fk;

	    next;
	}

	my $name = token(",");
	my %data;

	$tdata{null}{$name} = 1;
	
	while (@tokens) {
	    last if $tokens[0] eq ',';

	    if (match('CONSTRAINT')) {
		$data{constraint} = token(",");
	    }
	    elsif (match('DEFAULT')) {
		$tdata{default}{$name} = token(",");
	    }
	    elsif (match('NULL')) {
		$tdata{null}{$name} = 1;
	    }
	    elsif (match('NOT', 'NULL')) {
		$tdata{null}{$name} = 0;
	    }
	    elsif (match('PRIMARY', 'KEY')) {
		$tdata{pk}{name} = $data{constraint};
		$tdata{pk}{cols} = [ $name ];
		match('ASC') || match('DESC');
		match('AUTOINCREMENT');
	    }
	    elsif (match('UNIQUE')) {
		# skip
	    }
	    elsif (match('CHECK')) {
		# just skip it...
		token(",");
	    }
	    elsif (match('REFERENCES')) {

		my $fk = parse_refdef($table, [ $name ], $data{constraint})
		  or next;
		push @{$tdata{fk}}, $fk;
	    }
	    else {
		push @utok, token(",");
	    }

	}
	$tdata{type}{$name} = shift @utok;

	#printf "%s  : %s %s", $table, $name, $tdata{type}{$name};
	#printf " [%s]", $_ for @utok;
	#printf "\n";

    }
    continue {
	skipto(",");
	match(",");
    }

    $tables{$table} = \%tdata;
    
}

local *OUTPUT;
if ($opt_output) {
    open OUTPUT, ">", $opt_output
      or die "Error opening '$opt_output': $!\n";
}
else {
    *OUTPUT = *STDOUT;
}

for my $table (sort keys %tables) {
    my $def = $tables{$table};
    next unless $def->{fk};

    printf OUTPUT "--\n-- TABLE %s\n--\n", $table;

  FK:
    for my $fk (@{$def->{fk}}) {
	my $name = $fk->{name};
	my $tname = sprintf "%s_%s", $table, join("_", @{$fk->{srccols}});

	printf OUTPUT "\n-- FK %s.%s\n\n", $table, $name;

	unless (exists $tables{$fk->{reftable}}) {
	    error("unkown table `%s`", $fk->{reftable});
	    next FK if $opt_strict;	    
	}

	if ($fk->{refcols}) {
	    for my $col (@{$fk->{refcols}}) {
		unless (exists $tables{$fk->{reftable}}{type}{$col}) {
		    error("unknown column `%s`", $col);
		    next FK if $opt_strict;
		}
	    }
	}
	else {
	    $fk->{refcols} = $tables{$fk->{reftable}}{pk}{cols};
	}

	unless ($fk->{refcols}) {
	    error("missing reference columns");
	    next FK; # fatal
	}

	my $reftable = $fk->{reftable};
	my @refcols  = qname(@{$fk->{refcols}});
	my @srccols  = qname(@{$fk->{srccols}});

	# on INSERT

	printf OUTPUT qq{DROP TRIGGER IF EXISTS %s;\n}, qname("fki_$tname");
	printf OUTPUT qq{CREATE TRIGGER %s\n}, qname("${opt_prefix}fki_$tname");
	printf OUTPUT qq{BEFORE INSERT ON %s\n}, qname($table);
	printf OUTPUT qq{FOR EACH ROW BEGIN\n};
	printf OUTPUT qq{    SELECT RAISE(ROLLBACK, 'insert on table "%s" violates foreign key constraint "%s"')\n}, $table, $name;
	printf OUTPUT qq{    WHERE %s\n},
	  join(qq{\n      AND },
	       map { "NEW.$_ IS NOT NULL" } @srccols,
	      );
	printf(OUTPUT qq{      AND (SELECT %s FROM %s WHERE %s) IS NULL;\n},
	       join(", ", @refcols),
	       qname($reftable),
	       join(" AND ", 
		    map { sprintf "%s = NEW.%s", $refcols[$_], $srccols[$_] }
		    0..$#srccols
		   )
	      );
	printf OUTPUT "END;\n\n";

	# on UPDATE (forward)

	printf OUTPUT qq{DROP TRIGGER IF EXISTS %s;\n}, qname("fkuf_$tname");
	printf OUTPUT qq{CREATE TRIGGER %s\n}, qname("${opt_prefix}fkuf_$tname");
	printf(OUTPUT qq{BEFORE UPDATE OF %s ON %s\n},
	       join(", ", @srccols), qname($table),
	      );
	printf OUTPUT qq{FOR EACH ROW BEGIN\n};
	printf OUTPUT qq{    SELECT RAISE(ROLLBACK, 'update on table "%s" violates foreign key constraint "%s"')\n}, $table, $name;
	printf(OUTPUT qq{    WHERE %s\n},
	       join(qq{\n      AND },
		    map { "NEW.$_ IS NOT NULL" } @srccols,
		   ),
	      );
	printf(OUTPUT qq{      AND (SELECT %s FROM %s WHERE %s) IS NULL;\n},
	       join(", ", @refcols),
	       qname($reftable),
	       join(" AND ", 
		    map { sprintf "%s = NEW.%s", $refcols[$_], $srccols[$_] }
		    0..$#srccols
		   )
	      );
	printf OUTPUT "END;\n\n";

	# on UPDATE (reverse)

	if ($fk->{on}{update} eq 'error') {
	    printf OUTPUT qq{DROP TRIGGER IF EXISTS %s;\n}, qname("fkur_$tname");
	    printf OUTPUT qq{CREATE TRIGGER %s\n}, qname("${opt_prefix}fkur_$tname");
	    printf(OUTPUT qq{BEFORE UPDATE OF %s ON %s\n}, 
		   join(", ", @refcols), qname($reftable),
		  );
	    printf OUTPUT qq{FOR EACH ROW BEGIN\n};
	    printf OUTPUT qq{    SELECT RAISE(ROLLBACK, 'update on table "%s" violates foreign key constraint "%s"')\n}, $reftable, $name;
	    printf(OUTPUT qq{    WHERE (SELECT %s FROM %s WHERE %s) IS NOT NULL;\n},
		   join(", ", @srccols),
		   qname($table),
		   join(" AND ", 
			map { sprintf "%s = OLD.%s",
				$srccols[$_], $refcols[$_] }
		    0..$#srccols
		   )
	      );
	    printf OUTPUT "END;\n\n";
	}
	elsif ($fk->{on}{update} eq 'cascade') {
	    printf OUTPUT qq{DROP TRIGGER IF EXISTS %s;\n}, qname("fkur_$tname");
	    printf OUTPUT qq{CREATE TRIGGER %s\n}, qname("${opt_prefix}fkur_$tname");
	    printf(OUTPUT qq{AFTER UPDATE OF %s ON %s\n}, 
		   join(", ", @refcols), qname($reftable),
		  );
	    printf OUTPUT qq{FOR EACH ROW BEGIN\n};
	    printf(OUTPUT qq{    UPDATE %s SET %s WHERE %s;\n},
		   qname($table),
		   join(", ",
			map { sprintf "%s = NEW.%s",
				$srccols[$_], $refcols[$_] }
			0..$#srccols
		       ),
		   join(" AND ",
			map { sprintf "%s = OLD.%s",
				$srccols[$_], $refcols[$_] }
			0..$#srccols
		       ),
		  );
	    printf OUTPUT "END;\n\n";
	}
	elsif ($fk->{on}{update} eq 'null') {
	    printf OUTPUT qq{DROP TRIGGER IF EXISTS %s;\n}, qname("fku_$tname");
	    printf OUTPUT qq{CREATE TRIGGER %s\n}, qname("${opt_prefix}fku_$tname");
	    printf OUTPUT qq{BEFORE UPDATE ON %s\n}, qname($reftable);
	    printf OUTPUT qq{FOR EACH ROW BEGIN\n};
	    printf(OUTPUT qq{    UPDATE %s SET %s WHERE %s;\n},
		   qname($table),
		   join(", ", map { sprintf "%s = NULL", $_ } @srccols ),
		   join(" AND ",
			map { sprintf "%s = OLD.%s",
				$srccols[$_], $refcols[$_] }
			0..$#srccols
		       ),
		  );
	    printf OUTPUT "END;\n\n";
	}
	elsif ($fk->{on}{update} eq 'default') {
	    printf OUTPUT qq{DROP TRIGGER IF EXISTS %s;\n}, qname("fku_$tname");
	    printf OUTPUT qq{CREATE TRIGGER %s\n}, qname("${opt_prefix}fku_$tname");
	    printf OUTPUT qq{BEFORE UPDATE ON %s\n}, qname($reftable);
	    printf OUTPUT qq{FOR EACH ROW BEGIN\n};
	    printf(OUTPUT qq{    UPDATE %s SET %s WHERE %s;\n},
		   qname($table),
		   join(", ", 
			map { sprintf "%s = %s", $_, 
				$fk->{default}{$_}||'NULL' } @srccols,
		       ),
		   join(" AND ",
			map { sprintf "%s = OLD.%s", 
				$srccols[$_], $refcols[$_] }
			0..$#srccols
		       ),
		  );
	    printf OUTPUT "END;\n\n";
	}
	else {
	    error("Unsupported ON UPDATE: %s", $fk->{on}{update});
	}

	# on DELETE

	if ($fk->{on}{delete} eq 'error') {
	    printf OUTPUT qq{DROP TRIGGER IF EXISTS %s;\n}, qname("fkd_$tname");
	    printf OUTPUT qq{CREATE TRIGGER %s\n}, qname("${opt_prefix}fkd_$tname");
	    printf OUTPUT qq{BEFORE DELETE ON %s\n}, qname($reftable);
	    printf OUTPUT qq{FOR EACH ROW BEGIN\n};
	    printf OUTPUT qq{    SELECT RAISE(ROLLBACK, 'delete on table "%s" violates foreign key constraint "%s"')\n}, $reftable, $name;
	    printf(OUTPUT qq{    WHERE (SELECT %s FROM %s WHERE %s) IS NOT NULL;\n},
		   join(", ", @srccols),
		   qname($table),
		   join(" AND ", 
			map { sprintf "%s = OLD.%s",
				$srccols[$_], $refcols[$_] }
		    0..$#srccols
		   )
	      );
	    printf OUTPUT "END;\n\n";
	}
	elsif ($fk->{on}{delete} eq 'cascade') {
	    printf OUTPUT qq{DROP TRIGGER IF EXISTS %s;\n}, qname("fkd_$tname");
	    printf OUTPUT qq{CREATE TRIGGER %s\n}, qname("${opt_prefix}fkd_$tname");
	    printf OUTPUT qq{BEFORE DELETE ON %s\n}, qname($reftable);
	    printf OUTPUT qq{FOR EACH ROW BEGIN\n};
	    printf(OUTPUT qq{    DELETE FROM %s WHERE %s;\n},
		   qname($table),
		   join(" AND ",
			map { sprintf "%s = OLD.%s",
				$srccols[$_], $refcols[$_] }
			0..$#srccols
		       ),
		  );
	    printf OUTPUT "END;\n\n";
	}
	elsif ($fk->{on}{delete} eq 'null') {
	    printf OUTPUT qq{DROP TRIGGER IF EXISTS %s;\n}, qname("fkd_$tname");
	    printf OUTPUT qq{CREATE TRIGGER %s\n}, qname("${opt_prefix}fkd_$tname");
	    printf OUTPUT qq{BEFORE DELETE ON %s\n}, qname($reftable);
	    printf OUTPUT qq{FOR EACH ROW BEGIN\n};
	    printf(OUTPUT qq{    UPDATE %s SET %s WHERE %s;\n},
		   qname($table),
		   join(", ", map { sprintf "%s = NULL", $_ } @srccols ),
		   join(" AND ",
			map { sprintf "%s = OLD.%s",
				$srccols[$_], $refcols[$_] }
			0..$#srccols
		       ),
		  );
	    printf OUTPUT "END;\n\n";
	}
	elsif ($fk->{on}{delete} eq 'default') {
	    printf OUTPUT qq{DROP TRIGGER IF EXISTS %s;\n}, qname("fkd_$tname");
	    printf OUTPUT qq{CREATE TRIGGER %s\n}, qname("${opt_prefix}fkd_$tname");
	    printf OUTPUT qq{BEFORE DELETE ON %s\n}, qname($reftable);
	    printf OUTPUT qq{FOR EACH ROW BEGIN\n};
	    printf(OUTPUT qq{    UPDATE %s SET %s WHERE %s;\n},
		   qname($table),
		   join(", ", 
			map { sprintf "%s = %s", $_, 
				$fk->{default}{$_}||'NULL' } @srccols,
		       ),
		   join(" AND ",
			map { sprintf "%s = OLD.%s", 
				$srccols[$_], $refcols[$_] }
			0..$#srccols
		       ),
		  );
	    printf OUTPUT "END;\n\n";
	}
	else {
	    error("Unsupported ON DELETE: %s", $fk->{on}{delete});
	}
    }
}

#print Dumper \%tables;

exit 0;

# ---- helpers

sub parse_refdef {
    my($table, $srccols, $constraint) = @_;
    my %fk;

    $fk{reftable} = token(",") or return;
    $fk{srccols} = $srccols;
    
    if ($tokens[0] =~ /^\((.*)\)$/) {
	$fk{refcols} = [ split ",", $1 ];
	token();
	for (@{$fk{refcols}}) {
	    s/^\s+//;
	    s/\s+$//;
	}
    }
    
    # [ MATCH FULL | MATCH PARTIAL | MATCH SIMPLE ]
    if (match('MATCH')) {
	if (match('FULL')) {
	    $fk{match} = 'full';
	}
	elsif (match('PARTIAL')) {
	    $fk{match} = 'partial';
	}
	elsif (match('SIMPLE')) {
	    $fk{match} = 'simple';
	}
	else {
	    token(",");
	}
    }
    
    $fk{on} = { delete => 'error', 
		update => 'error',
	      };
    
    # [ ON DELETE action ] | [ ON UPDATE action ]
    while (match('ON')) {
	my($method, $action);
	
	if (match('DELETE')) {
	    $method = 'delete';
	}
	elsif (match('UPDATE')) {
	    $method = 'update';
	}
	else {
	    token(",") or return;
	}
	
	if (match('SET', 'NULL')) {
	    $action = 'null';
	}
	elsif (match('SET', 'DEFAULT')) {
	    $action = 'default';
	}
	elsif (match('CASCADE')) {
	    $action = 'cascade';
	}
	elsif (match('RESTRICT')) {
	    $action = 'error';
	}
	elsif (match('NO', 'ACTION')) {
	    $action = 'error'; # as restrict, but deferable
	}
	else {
	    token(",") or return;
	}

	$fk{on}{$method} = $action;
    }
    
    $fk{name} = $constraint ||
      sprintf "fk_%s_%s", $table, join("_", @{$fk{srccols}});
    
    return \%fk;
}

sub error {
    my($fmt,@args) = @_;
    printf "-- error: $fmt.\n", @args;
}

sub qname {
    my(@names) = @_;
    for (@names) {
	s/^`(.*)`$/$1/ or s/`/``/g;
	$_ = "`$_`";
    }
    return wantarray ? @names : $names[0];
}

sub match {
    for my $i (0..$#_) {
	return unless uc $tokens[$i] eq uc $_[$i];
    }
    return splice @tokens, 0, scalar @_;
}

sub token {
    for (@_) {
	return if uc $tokens[0] eq uc;
    }
    return shift @tokens;
}

sub skipto {
    shift @tokens while @tokens && uc $tokens[0] ne uc $_[0];
}

__END__

=head1 AUTHOR

Written by Peder Stray <peder.stray@gmail.com>

=head1 COPYRIGHT

Copyright E<169> 2008-2009 by Peder Stray <peder.stray@gmail.com>

Release under GPLv2 L<http://gnu.org/licenses/gpl.html> and distributed
from L<http://code.google.com/p/sqlite-fk-triggers/>

=cut
